import torch
from sklearn.metrics import f1_score

def binary_f1_score(scores, targets):
    """计算F1值"""
    y_true = targets.cpu().numpy()
    y_pred = scores.argmax(dim=1).cpu().numpy()
    return f1_score(y_true, y_pred, average='binary')

def train_epoch(model, optimizer, device, data_loader):
    """训练一个epoch"""
    model.train()
    iter = 0
    train_loss = 0
    train_f1 = 0
    for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
        batch_graphs = batch_graphs.to(device)
        batch_x = batch_graphs.ndata['feat'].to(device)  # num x feat
        batch_e = batch_graphs.edata['feat'].to(device)
        batch_labels = batch_labels.to(device)
        optimizer.zero_grad()

        batch_scores = model.forward(batch_graphs, batch_x, batch_e)
        loss = model.loss(batch_scores, batch_labels)
        loss.backward()
        optimizer.step()
        train_loss += loss.detach().item()
        train_f1 += binary_f1_score(batch_scores, batch_labels)
    train_loss /= (iter + 1)
    train_f1 /= (iter + 1)

    return train_loss, train_f1, optimizer


def evaluate_epoch(model, device, data_loader):
    """验证/测试一个epoch"""
    model.eval()
    iter = 0
    val_loss = 0
    val_f1 = 0
    with torch.no_grad():
        for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
            batch_graphs = batch_graphs.to(device)
            batch_x = batch_graphs.ndata['feat'].to(device)
            batch_e = batch_graphs.edata['feat'].to(device)
            batch_labels = batch_labels.to(device)

            batch_scores = model.forward(batch_graphs, batch_x, batch_e)
            loss = model.loss(batch_scores, batch_labels)
            val_loss += loss.detach().item()
            val_f1 += binary_f1_score(batch_scores, batch_labels)
        val_loss /= (iter + 1)
        val_f1 /= (iter + 1)

    return val_loss, val_f1

